import random
import numpy as np
from typing import List, Tuple, Optional, Union, Any
import time
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib.patches import RegularPolygon

plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号 '-' 显示为方块的问题


class WeightedGeneticPathFinder:
    def __init__(self, weight_map: List[List[int]], start: Tuple[int, int],
                 end: Tuple[int, int], population_size: int = 100,
                 max_generations: int = 500, mutation_rate: float = 0.1,
                 top_paths_count: int = 3):
        """
        初始化带权重的遗传算法路径查找器

        Args:
            weight_map: 二维权重地图数组，0-5表示通行难度，0最容易，5最难
            start: 起点坐标 (row, col)
            end: 终点坐标 (row, col)
            population_size: 种群大小
            max_generations: 最大迭代代数
            mutation_rate: 变异概率
            top_paths_count: 要追踪和显示的最优路径数量
        """
        self.top_paths = None
        self.weight_map = np.array(weight_map)
        self.rows, self.cols = self.weight_map.shape
        self.start = start
        self.end = end
        self.population_size = population_size
        self.max_generations = max_generations
        self.mutation_rate = mutation_rate
        self.top_paths_count = top_paths_count  # 新增参数，控制追踪的路径数量

        # 验证起点和终点是否有效
        if not self._is_valid_position(start) or not self._is_valid_position(end):
            raise ValueError("起点或终点位于不可通行区域")

    def _is_valid_position(self, pos: Tuple[int, int]) -> bool:
        """检查位置是否有效（权重不为6，即不是障碍物）"""
        row, col = pos
        return (0 <= row < self.rows and
                0 <= col < self.cols and
                self.weight_map[row][col] < 6)  # 6表示障碍物

    def _get_position_weight(self, pos: Tuple[int, int]) -> int:
        """获取位置的权重值"""
        row, col = pos
        return int(self.weight_map[row][col])

    def _generate_random_path(self) -> List[Tuple[int, int]]:
        """生成偏向终点的随机路径"""
        path = [self.start]
        current = self.start

        max_steps = self.rows * self.cols * 2
        steps = 0

        while current != self.end and steps < max_steps:
            neighbors = self._get_neighbors(current, path)

            if not neighbors:
                if len(path) > 1:
                    path.pop()
                    current = path[-1]
                else:
                    break
                continue

            # 偏向选择靠近终点的方向
            best_neighbor = None
            best_distance = float('inf')

            for neighbor in neighbors:
                distance = abs(neighbor[0] - self.end[0]) + abs(neighbor[1] - self.end[1])
                if distance < best_distance:
                    best_distance = distance
                    best_neighbor = neighbor

            # 80%概率选择最佳方向，20%随机选择
            if random.random() < 0.8 and best_neighbor:
                current = best_neighbor
            else:
                current = random.choice(neighbors)

            path.append(current)
            steps += 1

        return path

    def _get_neighbors(self, current: Union[Tuple[int, int], Any], path: List[Tuple[int, int]]) -> List[Any]:
        neighbors = []
        row, col = current

        # 仅保留上下左右四个方向的相邻节点
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
        for dr, dc in directions:
            new_row, new_col = row + dr, col + dc
            new_pos = (new_row, new_col)

            # 检查新位置是否有效且未访问过
            if (self._is_valid_position(new_pos) and
                    new_pos not in path):
                neighbors.append(new_pos)
        return neighbors

    def calculate_path_cost(self, path: List[Tuple[int, int]]) -> float:
        """计算路径总成本（权重和）"""
        if not path:
            return float('inf')

        path_total_cost = 0
        for pos in path:
            path_total_cost += self._get_position_weight(pos)
        return path_total_cost

    def _calculate_fitness(self, path: List[Tuple[int, int]]) -> float:
        if not path or path[0] != self.start:
            return 0

        # 如果路径到达终点，给予高奖励
        if path[-1] == self.end:
            cost = self.calculate_path_cost(path)
            # 到达终点的路径应该获得比任何未到达路径更高的适应度
            base_reward = 1000  # 基础奖励
            return base_reward / (cost + 1)

        # 对于未到达终点的路径，主要奖励距离终点的接近程度
        last_pos = path[-1]
        manhattan_distance = abs(last_pos[0] - self.end[0]) + abs(last_pos[1] - self.end[1])

        # 减少对路径长度和成本的惩罚，重点奖励接近终点
        return 1.0 / (manhattan_distance + 1)  # 简化适应度计算

    def _crossover(self, parent1: List[Tuple[int, int]],
                   parent2: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
        """路径交叉操作"""
        # 简单的单点交叉
        if len(parent1) < 3 or len(parent2) < 3:
            return parent1[:] if random.random() < 0.5 else parent2[:]

        # 找到两条路径的公共节点
        common_nodes = list(set(parent1) & set(parent2))

        if common_nodes and len(common_nodes) > 1:
            # 在公共节点处分割并重组
            common_node = random.choice(common_nodes[1:])  # 不选择起点
            idx1 = parent1.index(common_node)
            idx2 = parent2.index(common_node)

            child = parent1[:idx1] + parent2[idx2:]
        else:
            # 随机交叉点
            cut_point = random.randint(1, min(len(parent1), len(parent2)) - 1)
            child = parent1[:cut_point] + parent2[cut_point:]

        # 修复路径中的冲突（重复节点）
        fixed_path = []
        visited = set()

        for node in child:
            if node not in visited:
                fixed_path.append(node)
                visited.add(node)

        # 如果末尾不是终点，尝试连接到终点
        if fixed_path and fixed_path[-1] != self.end:
            # 简单的尝试添加一步到终点
            last_node = fixed_path[-1]
            end_row, end_col = self.end
            last_row, last_col = last_node

            # 尝试直接向终点移动
            if abs(last_row - end_row) + abs(last_col - end_col) == 1:
                if self._is_valid_position(self.end):
                    fixed_path.append(self.end)

        return fixed_path

    def _mutate(self, path: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
        """路径变异操作"""
        if len(path) < 3 or random.random() > self.mutation_rate:
            return path[:]

        mutated_path = path[:]

        # 随机选择变异方式
        mutation_type = random.choice(['remove', 'insert', 'swap'])

        if mutation_type == 'remove' and len(mutated_path) > 2:
            # 随机移除一个中间节点
            idx = random.randint(1, len(mutated_path) - 2)
            mutated_path.pop(idx)
        elif mutation_type == 'insert' and len(mutated_path) > 1:
            # 在随机位置插入一个新的邻近节点
            idx = random.randint(1, len(mutated_path) - 1)
            prev_node = mutated_path[idx - 1]

            # 获取相邻可通行节点
            neighbors = self._get_neighbors(prev_node, path)

            if neighbors:
                new_node = random.choice(neighbors)
                mutated_path.insert(idx, new_node)
        elif mutation_type == 'swap' and len(mutated_path) > 3:
            # 交换两个中间节点
            idx1 = random.randint(1, len(mutated_path) - 2)
            idx2 = random.randint(1, len(mutated_path) - 2)
            if idx1 != idx2:
                mutated_path[idx1], mutated_path[idx2] = mutated_path[idx2], mutated_path[idx1]

        return mutated_path

    def _repair_path(self, path: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
        """尝试修复路径，使其到达终点"""
        if not path or path[-1] == self.end:
            return path

        current = path[-1]
        repaired_path = path[:]

        # 尝试从当前位置找到到终点的路径
        max_repair_steps = 50
        steps = 0

        while current != self.end and steps < max_repair_steps:
            # 选择最接近终点的邻居
            neighbors = self._get_neighbors(current, repaired_path)
            if not neighbors:
                break

            best_neighbor = min(neighbors,
                                key=lambda n: abs(n[0] - self.end[0]) + abs(n[1] - self.end[1]))

            repaired_path.append(best_neighbor)
            current = best_neighbor
            steps += 1

        return repaired_path

    def find_path(self) -> Optional[List[Tuple[int, int]]]:
        """执行遗传算法寻找路径"""
        # 记录开始时间
        start_time = time.time()

        # 初始化种群
        population = []
        for _ in range(self.population_size):
            individual = self._generate_random_path()
            population.append(individual)

        best_path = None
        best_fitness = 0
        
        # 添加存储前n路径的列表
        top_paths = []  # 存储 (path, fitness, cost) 元组

        for generation in range(self.max_generations):
            # 计算适应度
            fitness_scores = []
            for individual in population:
                fitness = self._calculate_fitness(individual)
                fitness_scores.append(fitness)

                # 更新最佳路径
                if fitness > best_fitness:
                    best_fitness = fitness
                    best_path = individual[:]

                # 更新前n路径列表
                path_cost = self.calculate_path_cost(individual) if individual[-1] == self.end else float('inf')
                top_paths.append((individual[:], fitness, path_cost))
                # 按适应度排序并保留前n个
                top_paths.sort(key=lambda x: x[1], reverse=True)
                top_paths = top_paths[:self.top_paths_count]

            # 如果找到完美解则提前结束
            if best_path and best_path[-1] == self.end:
                print(f"在第 {generation + 1} 代找到解决方案")
                # print(f"最佳路径成本: {best_cost}")
                # break

            # 选择新一代个体
            new_population = []

            # 保留最佳个体（精英主义）
            elite_count = self.population_size // 10
            sorted_indices = np.argsort(fitness_scores)[::-1]
            for i in range(elite_count):
                idx = sorted_indices[i]
                new_population.append(population[idx][:])

            # 填充剩余个体
            while len(new_population) < self.population_size:
                # 轮盘赌选择两个父代
                parent1 = self._roulette_selection(population, fitness_scores)
                parent2 = self._roulette_selection(population, fitness_scores)

                # 交叉生成子代
                child = self._crossover(parent1, parent2)

                # 变异
                child = self._mutate(child)

                new_population.append(child)

            population = new_population

            # 输出进度
            if generation % 50 == 0:
                print(f"第 {generation} 代，当前最佳适应度: {best_fitness}")

        # 计算执行时间
        end_time = time.time()
        execution_time = end_time - start_time

        # 输出执行时间信息
        print(f"算法执行时间: {execution_time:.4f} 秒")
        
        # 输出前n路径信息
        print(f"\n适应度前{self.top_paths_count}的路径:")
        for i, (path, fitness, cost) in enumerate(top_paths):
            print(f"第{i+1}名 - 适应度: {fitness:.6f}", end="")
            if path[-1] == self.end:
                print(f", 路径成本: {cost:.2f}, 路径: {path} (到达终点)")
            else:
                print(f", 路径: {path} (未到达终点)")

        # 保存前n路径以供可视化
        self.top_paths = top_paths

        return best_path if best_path and best_path[-1] == self.end else None

    @staticmethod
    def _roulette_selection(population: List[List[Tuple[int, int]]],
                            fitness_scores: List[float]) -> List[Tuple[int, int]]:
        """轮盘赌选择"""
        total_fitness = sum(fitness_scores)
        if total_fitness == 0:
            return random.choice(population)[:]

        pick = random.uniform(0, total_fitness)
        current = 0

        for i, fitness in enumerate(fitness_scores):
            current += fitness
            if current >= pick:
                return population[i][:]

        return population[-1][:]

    def visualize_map_and_path(self):
        """
        可视化权重地图和路径
        """
        # 创建图形和轴
        fig, ax = plt.subplots(figsize=(12, 10))

        # 创建颜色映射（0-5权重对应不同颜色，6为障碍物）
        colors = ['lightgreen', 'yellow', 'orange', 'red', 'darkred', 'purple', 'black']
        cmap = ListedColormap(colors)

        # 显示地图
        ax.imshow(self.weight_map, cmap=cmap, origin='upper')

        # 设置标题
        ax.set_title('权重地图和路径可视化')

        # 设置坐标轴标签
        ax.set_xlabel('列 (Column)')
        ax.set_ylabel('行 (Row)')

        # 添加网格线
        ax.set_xticks(np.arange(-0.5, self.cols, 1), minor=True)
        ax.set_yticks(np.arange(-0.5, self.rows, 1), minor=True)
        ax.grid(which='minor', color='gray', linestyle='-', linewidth=0.5)

        # 添加颜色条
        cbar = plt.colorbar(ax.images[0], ax=ax, shrink=0.8)
        cbar.set_label('通行权重 (0-5)，黑色为障碍物')

        # 标记起点和终点
        start_row, start_col = self.start
        end_row, end_col = self.end

        # 在起点画绿色圆圈
        ax.add_patch(plt.Circle((start_col, start_row), 0.3,
                                facecolor='lime', alpha=0.8, edgecolor='black', linewidth=2))
        ax.text(start_col, start_row, 'S', ha='center', va='center',
                fontsize=12, fontweight='bold', color='white')

        # 在终点画红色五角星
        ax.add_patch(RegularPolygon((end_col, end_row), 5, radius=0.4,
                                    facecolor='cyan', alpha=0.8, edgecolor='black', linewidth=2))
        ax.text(end_col, end_row, 'E', ha='center', va='center',
                fontsize=12, fontweight='bold', color='black')

        # 如果没有提供特定路径但有top_paths数据，则绘制前n路径
        if hasattr(self, 'top_paths') and self.top_paths:
            # 为不同路径生成颜色
            colormap = plt.colormaps.get_cmap('tab10')
            path_colors = [colormap(i) for i in range(min(len(self.top_paths), 10))]
            
            # 绘制前n路径
            for i, (top_path, fitness, cost) in enumerate(self.top_paths):
                if i >= self.top_paths_count:
                    break

                if top_path:
                    # 提取路径的行和列坐标
                    path_rows = [pos[0] for pos in top_path]
                    path_cols = [pos[1] for pos in top_path]

                    # 绘制路径线
                    ax.plot(path_cols, path_rows, color=path_colors[i % len(path_colors)], linewidth=3, marker='o',
                            markersize=6, markerfacecolor=path_colors[i % len(path_colors)], markeredgecolor='white', markeredgewidth=1,
                            label=f'第{i+1}名 (适应度: {fitness:.4f})')

                    # 在路径点上添加序号
                    for j, (row, col) in enumerate(top_path):
                        ax.text(col, row, str(j), ha='center', va='center',
                                fontsize=8, color='white', fontweight='bold',
                                bbox=dict(boxstyle="circle,pad=0.1", facecolor=path_colors[i % len(path_colors)], alpha=0.7))
            
            # 添加图例
            ax.legend(loc='upper right', bbox_to_anchor=(1, 1))

        # 设置坐标轴范围
        ax.set_xlim(-0.5, self.cols - 0.5)
        ax.set_ylim(self.rows - 0.5, -0.5)  # 翻转y轴使(0,0)在左上角

        # 显示图形
        plt.tight_layout()
        plt.show()


# 使用示例
if __name__ == "__main__":
    # 创建20x20的权重地图 (0-5表示通行难度，数值越大越难通行，6表示障碍物)
    weight_map_example = [
        [0, 1, 2, 6, 1, 0, 1, 2, 1, 4, 6, 1, 0, 2, 3, 1, 6, 1, 1, 0],
        [1, 6, 1, 6, 2, 6, 1, 1, 6, 2, 6, 2, 6, 1, 6, 2, 6, 1, 6, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [6, 6, 1, 6, 6, 1, 1, 6, 6, 1, 6, 6, 1, 6, 6, 1, 6, 1, 6, 6],
        [2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2],
        [1, 6, 6, 6, 1, 0, 1, 6, 6, 6, 1, 6, 6, 6, 1, 6, 6, 1, 6, 1],
        [0, 1, 2, 6, 1, 0, 1, 2, 1, 4, 6, 1, 0, 2, 3, 1, 6, 1, 1, 0],
        [1, 6, 1, 6, 2, 6, 1, 1, 6, 2, 6, 2, 6, 1, 6, 2, 6, 1, 6, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [6, 6, 1, 6, 6, 1, 1, 6, 6, 1, 6, 6, 1, 6, 6, 1, 6, 1, 6, 6],
        [2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2],
        [1, 6, 6, 6, 1, 0, 1, 6, 6, 6, 1, 6, 6, 6, 1, 6, 6, 1, 6, 1],
        [0, 1, 2, 6, 1, 0, 1, 2, 1, 4, 6, 1, 0, 2, 3, 1, 6, 1, 1, 0],
        [1, 6, 1, 6, 2, 6, 1, 1, 6, 2, 6, 2, 6, 1, 6, 2, 6, 1, 6, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [6, 6, 1, 6, 6, 1, 1, 6, 6, 1, 6, 6, 1, 6, 6, 1, 6, 1, 6, 6],
        [2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2],
        [1, 6, 6, 6, 1, 0, 1, 6, 6, 6, 1, 6, 6, 6, 1, 6, 6, 1, 6, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [0, 6, 6, 6, 0, 6, 1, 6, 0, 6, 6, 6, 0, 6, 6, 6, 0, 1, 6, 0]
    ]

    # 设置起点和终点
    start_point = (0, 0)
    end_point = (19, 19)

    # 创建路径查找器，指定要追踪和显示5条最优路径
    path_finder = WeightedGeneticPathFinder(
        weight_map=weight_map_example,
        start=start_point,
        end=end_point,
        population_size=50,
        max_generations=10,
        mutation_rate=0.1,
        top_paths_count=5  # 指定要追踪和显示5条最优路径
    )

    # 查找路径
    result_path = path_finder.find_path()

    if result_path:
        print("找到最优路径:")
        print(result_path)
        # 计算路径总成本
        total_cost = path_finder.calculate_path_cost(result_path)
        print(f"路径总成本: {total_cost}")
    else:
        print("未能找到有效路径")
    
    # 可视化前n路径
    path_finder.visualize_map_and_path()